import requests
import json
from typing import List, Dict, Optional
import time
import re

class ChapterNode:
    """章节树节点"""
    def __init__(self, title: str, content: str = "", level: int = 1):
        self.title = title  # 章节标题
        self.content = content  # 章节内容
        self.level = level  # 章节层级(1: 主章节, 2: 二级章节, 3: 三级章节)
        self.children = []  # 子章节列表
        self.full_title = title  # 完整标题路径
        
    def add_child(self, child: 'ChapterNode'):
        # 更新子节点的完整标题
        child.full_title = f"{self.full_title} > {child.title}"
        self.children.append(child)
        
    def get_full_content(self) -> str:
        """获取包含所有子章节的完整内容"""
        # 使用完整标题替换原标题
        result = [f"{self.full_title}\n{self.content}"]
        for child in self.children:
            result.append(child.get_full_content())
        return '\n'.join(result)

class QAGenerator:
    """问答对生成器 - 专门针对操作手册文档优化"""
    
    def __init__(self, config: dict):
        """初始化配置"""
        self.config = config
        self.api_key = config['ai']['api_key']
        self.api_base = config['ai']['api_base']
        self.model = config['ai']['model']
        self.temperature = config['ai']['temperature']
        self.max_tokens = config['ai']['max_tokens']


    def generate(self, content: str) -> Optional[List[Dict]]:
        """生成问答对"""
        sections = self._split_content(content)
        all_qa_pairs = []
        
        for section in sections:
            if self._is_invalid_section(section):
                continue
                
            prompt = self._create_clean_contents_prompt(section)
            
            try:
                # 第一次生成QA对
                response = self._call_api('开始吧', prompt, require_json=True)
                if 'choices' in response and response['choices']:
                    qa_pairs = []
                    try:
                        qa_pairs = json.loads(response['choices'][0]['message']['content'])
                        print(f"\n\n\n优化前的Q&A：{qa_pairs} ")
                        # 传入原文内容进行优化
                        origin_content = response['choices'][0]['message']['content']
                        qa_pairs = self._optimize_qa_pairs(origin_content, section)
                        
                        # 转换图片链接格式
                        qa_pairs = self._convert_image_format(qa_pairs)

                    except Exception as e:
                        print(f"提示词是：{prompt}"
                              f"响应出现：{response['choices'][0]['message']['content']}"
                              f"\n"
                              f"问题：{str(e)}")

                    valid_pairs = [qa for qa in qa_pairs if qa['A'].strip()]
                    all_qa_pairs.extend(valid_pairs)
            except Exception as e:
                print(f"处理章节时出错: {str(e)}")
                continue
            
        return all_qa_pairs

    def _split_content(self, content: str) -> List[str]:
        """将文档按章节分割并构建章节树结构
        
        Args:
            content: markdown格式的文档内容
        Returns:
            按长度限制组织好的文档片段列表
        """
        # 解析所有章节
        content = '\n'.join(filter(lambda x: x.strip(), content.split('\n')))
        chunks = self._parse_chapters(content)
        
        # 构建章节树
        chapter_tree = self._build_chapter_tree(chunks)
        
        # 基于章节树生成合适大小的内容片段
        return self._generate_content_chunks(chapter_tree)

    def _parse_chapters(self, content: str) -> List[Dict]:
        """解析文档中的所有章节
        
        Returns:
            List[Dict]: 包含章节信息的字典列表
            [{"title": "1 引言", "content": "...", "level": 1}, ...]
        """
        chunks = []
        current_chunk = []
        
        # 章节标题匹配模式
        main_chapter_pattern = re.compile(r'^(\d+)\s+[^\d]')  # 1 引言
        sub_chapter_pattern = re.compile(r'^(\d+\.\d+)')      # 1.1 xxx
        sub_sub_chapter_pattern = re.compile(r'^(\d+\.\d+\.\d+)') # 1.1.1 xxx
        
        lines = content.split('\n')
        
        for line in lines:
            # 判断行是否是章节标题
            main_match = main_chapter_pattern.match(line)
            sub_match = sub_chapter_pattern.match(line)
            sub_sub_match = sub_sub_chapter_pattern.match(line)
            
            is_new_chapter = any([main_match, sub_match, sub_sub_match])
            
            if is_new_chapter:
                # 保存当前章节
                if current_chunk:
                    chunks.append({
                        "title": current_chunk[0],
                        "content": '\n'.join(current_chunk),
                        "level": self._get_chapter_level(current_chunk[0])
                    })
                current_chunk = [line]
            else:
                current_chunk.append(line)
        
        # 添加最后一个章节
        if current_chunk:
            chunks.append({
                "title": current_chunk[0],
                "content": '\n'.join(current_chunk),
                "level": self._get_chapter_level(current_chunk[0])
            })
        
        return chunks

    def _get_chapter_level(self, title: str) -> int:
        """获取章节层级"""
        if re.match(r'^\d+\.\d+\.\d+', title):
            return 3
        elif re.match(r'^\d+\.\d+', title):
            return 2
        elif re.match(r'^\d+\s+', title):
            return 1
        return 0

    def _build_chapter_tree(self, chunks: List[Dict]) -> List[ChapterNode]:
        """构建章节树"""
        root_chapters = []  # 存储主章节
        current_l1 = None  # 当前一级章节
        current_l2 = None  # 当前二级章节
        
        for chunk in chunks:
            node = ChapterNode(chunk['title'], chunk['content'], chunk['level'])
            
            if chunk['level'] == 1:
                current_l1 = node
                root_chapters.append(node)
                current_l2 = None
            elif chunk['level'] == 2 and current_l1:
                current_l2 = node
                current_l1.add_child(node)  # 这里会自动设置完整标题
            elif chunk['level'] == 3 and current_l2:
                current_l2.add_child(node)  # 这里会自动设置完整标题
        
        return root_chapters

    def _generate_content_chunks(self, chapter_tree: List[ChapterNode]) -> List[str]:
        """基于章节树生成合适大小的内容片段"""
        chunks = []
        max_length = 2000  # API输入的最大长度限制
        
        for chapter in chapter_tree:
            # 获取主章节的完整内容并去除多余空行
            full_content = chapter.get_full_content()

            if len(full_content) <= max_length:
                # 如果主章节内容不超长，作为一个整体
                chunks.append(full_content)
            else:
                # 如果主章节内容超长，按二级章节拆分
                for child in chapter.children:
                    child_content = child.get_full_content()
                    if len(child_content) <= max_length:
                        chunks.append(child_content)
                    else:
                        # 如果二级章节还是太长，可以继续拆分或者截断
                        for grandchild in child.children:
                            chunks.append(grandchild.get_full_content())
        
        return chunks

    def _create_clean_contents_prompt(self, content: str) -> str:
        """创建清洗后的内容提示词"""
        # 移除页码标记
        content = re.sub(r'## 第\d+页', '', content)
        
        # 移除目录样式的点线
        content = re.sub(r'\.{3,}.*?\d+', '', content)
        
        # 移除重复的标题
        lines = content.split('\n')
        clean_lines = []
        prev_line = None
        for line in lines:
            if line != prev_line:
                clean_lines.append(line)
            prev_line = line
        content = '\n'.join(clean_lines)
        
        # 构建提示词
        prompt = """请根据以下【内容】生成实用的json数组问答对。请遵守以下规则：

1. 【内容】筛选:
   - 只提取具体的操作步骤、功能说明、界面说明等实际应用内容
   - 忽略目录、版权声明、文档版本、操作文档描述、目的等非实质内容
   - 如果内容是纯标题、目录或无实际操作指导的内容，返回空数组

2. 问答对格式:
   - 问题必须包含完整的模块路径和具体操作，例如:
     * "在[模块路径]中，如何[具体操作]?"
     * "在[功能位置]中，[功能名称]的作用是什么?"
     * "使用[具体功能]时，如何解决[具体问题]?"
   - 问题示例：
     * "在配置中心的品质配置-不良管理模块中，如何新增不良类别？"
     * "在Web端的不良类别管理界面中，如何进行批量导入？"
     * "在Proline端的数据采集界面中，如何设置产品不良？"
   - 答案要完整保留原文的操作步骤和图片
   - 答案不要出现编造，内容要与原文一致
   - 使用</br>表示换行
   - 严格按照 ![image](url)内容处理

3. ![image](url)处理规则（最重要）：
   - 视作图片链接，需要完整保留在回答中
   - 禁止修改链接，禁止给出不属于原文的链接。
   - 连续的 ![image](url)只有空格或空行的多张图片，视作一张图连续展示，禁止丢失其一
   - 图片链接所在内容位置，需保持和原文一致

    

4. 问题生成规则：
   - 每个问题必须包含完整的定位信息：
     * 系统平台（如Web端、Proline端）
     * 功能路径（如配置中心->品质配置->不良管理）
     * 具体模块（如不良类别管理、工序不良表等）
   
   - 问题应该是完整的句子，避免使用简单的关键词
   - 问题应该具有唯一性，不同模块的相似操作要通过模块路径区分
   - 不生成功能或界面包含哪些部分的问题
5. 返回格式:
[
  {
    "Q": "完整的问题描述，包含模块路径",
    "A": "详细的答案，包含完整步骤和图片"
  }
]
6.JSON格式要求：
   - 使用标准JSON格式，不要包含任何转义字符或特殊字符
   - 换行使用<br>标记，不使用\\n或\\r\\n
   - 不要使用制表符\\t或其他空白控制字符
   - 答案内容应该是连续的文本，避免不必要的空格和缩进

【内容】：
"""
        return prompt + content

    def _is_invalid_section(self, section: str) -> bool:
        """判断章节是否为无效内容"""
        # 包含大量点线的可能是目录
        if '.' * 10 in section:
            return True
        
        # 只包含标题的内容
        if len(section.split('\n')) <= 3:
            return True
        
        # 不包含实际内容的章节(如纯数字标题)
        if re.match(r'^\d+(\.\d+)*\s+\w+\s*$', section.strip()):
            return True
        
        # 其他无效内容判断...
        invalid_keywords = ['目录', '版权声明', '文档版本', '编制人', '审核人']
        if any(keyword in section for keyword in invalid_keywords):
            return True
        
        return False

    def _call_api(self, prompt: str, system_prompt: str = None, history: List[Dict] = None, require_json: bool = False) -> dict:
        """调用AI API - 使用百炼API
        
        Args:
            prompt: 用户提示词
            system_prompt: 系统角色设定
            history: 历史对话记录 [{"role": "user/assistant", "content": "xxx"}]
            require_json: 是否要求返回JSON格式
        """
        # 重试机制
        max_retries = 3
        for attempt in range(max_retries):
            try:
                # 准备消息列表
                messages = []
                
                # 添加系统提示词
                if system_prompt:
                    messages.append({
                        'role': 'system',
                        'content': system_prompt
                    })
                
                # 添加历史对话
                if history:
                    messages.extend(history)
                    
                # 添加当前提示词
                messages.append({
                    'role': 'user',
                    'content': prompt
                })
                
                # 准备请求数据
                headers = {
                    'Content-Type': 'application/json',
                    'Authorization': f'Bearer {self.api_key}'
                }
                
                data = {
                    'model': 'qwen2.5-7b-instruct-1m',  # 使用通义千问plus模型
                    'messages': messages,
                    'temperature': self.temperature,
                    'stream': False,
                }
                
                # 仅在需要JSON返回时添加response_format
                if require_json:
                    data['response_format'] = {'type': 'json_object'}
                
                # 发送请求
                response = requests.post(
                    "https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions",
                    headers=headers,
                    json=data,
                    timeout=3000
                )
                
                # 检查响应
                if response.status_code == 200:
                    return response.json()
                else:
                    print(f"百炼API调用失败: {response.status_code} - {response.text}")
                    if attempt == max_retries - 1:
                        print(f"提示词：{system_prompt}")
                        raise Exception(f"API调用失败: {response.status_code}")
                    time.sleep(1)  # 失败后等待1秒重试
                    
            except Exception as e:
                print(f"API调用异常: {str(e)}")
                print(f"提示词：{system_prompt}")
                if attempt == max_retries - 1:
                    raise
                time.sleep(1)  # 异常后等待1秒重试
                
            return {}

    def _optimize_qa_pairs(self, qa_pairs: str, original_content: str) -> List[Dict]:
        """对生成的QA对进行二次优化，确保与原文内容一致
        
        Args:
            qa_pairs: 初次生成的QA对列表
            original_content: 原始文档内容
        """
        optimization_prompt = """根据【原文内容】对【生成的问答对】进行严格审核，修复回复中不符合【要求】，进行纠正返回：
##要求
1.针对【生成的问答对】中Q也就是问题的准确性：
    - 确保问题包含完整的功能定位（平台->模块->功能）
    - 问题必须对应原文中实际存在的操作或说明
    - 删除无法从原文中找到明确答案的问题
2.针对【生成的问答对】中A也就是回答的准确性:
    - 确保【生成的问答对】内容需要原文内容出现明确说明，表述方式一致。
    - 保持内容的顺序
    - 保持其中图片链接和原文顺序一致
    - 将<img src="url"/> 方式展示图片方式改为 ![image](url)
    - 再次强调不得生成内容没有的内容和图片链接
    - 原文中保持连续多张的![image](url)，若是丢失，需要恢复补充丢失链接。
    - 禁止根据问题编造原文没有的答案
3.JSON格式要求：
   - 使用标准JSON格式，不要包含任何转义字符或特殊字符
   - 换行使用<br>标记，禁止使用\\n或\\r\\n
   - 禁止使用制表符\\t或其他空白控制字符
   - 答案内容应该是连续的文本，避免不必要的空格和缩进

【原文内容】：
"""





        try:
            # 准备优化提示词
            prompt = optimization_prompt + original_content + """\n【生成的问答对】：""" + f"{qa_pairs}" +"请返回优化后的JSON格式的Q&A列表"

            # 调用API进行优化
            response = self._call_api('开始优化', prompt, require_json=True)
            
            if 'choices' in response and response['choices']:
                try:
                    optimized_qa = json.loads(response['choices'][0]['message']['content'])
                    # 验证优化后的QA对
                    if self._validate_qa_pairs(optimized_qa, original_content):
                        print(f"优化后的Q&A：{qa_pairs}")
                        print(f'优化提示词：{prompt}')
                        return optimized_qa
                except json.JSONDecodeError as e:
                    print(f"优化后的QA对解析失败: {str(e)}")
                    return qa_pairs
            print(f'优化提示词：{prompt}')
        except Exception as e:
            print(f"优化QA对时出错: {str(e)}")
            return qa_pairs
            
        return qa_pairs

    def _validate_qa_pairs(self, qa_pairs: List[Dict], original_content: str) -> bool:
        """验证优化后的QA对是否符合要求
        
        Args:
            qa_pairs: 优化后的QA对列表
            original_content: 原始文档内容
        
        Returns:
            bool: 验证是否通过
        """
        try:
            for qa in qa_pairs:
                # 检查必要字段
                if not all(key in qa for key in ('Q', 'A')):
                    return False
                
                # 检查问题格式
                if not ('如何' in qa['Q'] or '什么' in qa['Q'] or '怎样' in qa['Q']):
                    return False
                
                # 检查答案格式
                if '<br>' not in qa['A']:
                    return False
                
                # 检查图片链接格式
                if '![image]' in qa['A']:
                    if not re.search(r'!\[image\]\([^)]+\)', qa['A']):
                        return False
                
            return True
            
        except Exception as e:
            print(f"QA对验证失败: {str(e)}")
            return False

    def _convert_image_format(self, qa_pairs: List[Dict]) -> List[Dict]:
        """转换QA对中的图片链接格式
        
        将 <image src="url"> 格式转换为 ![图片](url) 格式
        """
        if not qa_pairs:
            return qa_pairs
            
        for qa in qa_pairs:
            if 'A' in qa:
                # 匹配 <image src="url"> 或 <img src="url"> 格式
                qa['A'] = re.sub(
                    r'<im(?:age|g)\s+src=["\']([^"\']+)["\'](?:\s*/)?>', 
                    r'![图片](\1)', 
                    qa['A']
                )
                
                # 匹配 <image src=url> 或 <img src=url> 格式（无引号）
                qa['A'] = re.sub(
                    r'<im(?:age|g)\s+src=([^\s>]+)(?:\s*/)?>', 
                    r'![图片](\1)', 
                    qa['A']
                )
                
        return qa_pairs
if __name__ == '__main__':
    text = """在Web端中，如果遇到页面内容空白的问题，请按照以下步骤操作：<br><ol><li>检查浏览器地址栏，确认是否使用https访问系统。操作方法：复制地址栏的网址，粘贴到记事本中，查看是否为https://xxxx。<br><img src="https://ums-operation.obs.cn-north-4.myhuaweicloud.com/testDoc/images/1739329067_340f1587.jpeg" alt="图 4.1 查看访问网址"></li><li>如果是https访问系统，尝试清除浏览器缓存，特别是【自动填充表单数据】。以Chrome浏览器为例，具体操作如下：</li></ol><ol start="2"><li>点击【隐私和安全】，再点击【清除浏览数据】，打开【清除浏览数据】弹窗。<br><img src="https://ums-operation.obs.cn-north-4.myhuaweicloud.com/testDoc/images/1739329067_9fa7b6e0.jpeg" alt="图 4.2 打开清除浏览数据"></li><li>注意勾选【自动填充表单数据】，然后点击【清除数据】，清除缓存。<br><img src="https://ums-operation.obs.cn-north-4.myhuaweicloud.com/testDoc/images/1739329067_3f008058.jpeg" alt="图4.3 清除数据"><br><img src="https://ums-operation.obs.cn-north-4.myhuaweicloud.com/testDoc/images/1739329068_171a6695.jpeg" alt=""></li><li>清除缓存后，使用http://xxx访问系统，若页面数据仍为空白，请联系管理员。</li></ol>"""
    pattern = r'<img[^<]+<url[^>]*>([^<]+)</url>.*?alt="([^"]+)"'


    # 替换为 ![AltText](URL) 格式

    pattern = r'<img src="(.+?)" alt="(.+?)">'


    def replace_img(match):
        url = match.group(1)
        alt = match.group(2)
        return f'![{alt}]({url})'
    # 使用正则表达式进行替换
    converted_text = re.sub(pattern, replace_img, text)
    print(converted_text)




